"""
指数分析智能体模块 - 提供指数分析功能
"""

from typing import Dict, TypedDict, List, Union, Any, Callable
from langgraph.graph import StateGraph, END
import json
import datetime
import logging
import pandas as pd

from FDAS.agents.base import BaseAgent, BaseAgentState
from FDAS.data import get_asset
from FDAS.data.assets.base import ASSET_TYPE_INDEX
from FDAS.agents.workflows.prompts import (
    get_query_parse_prompt,
    get_specific_analysis_prompt,
    get_synthesis_prompt
)

# 自定义JSON编码器，处理日期类型
class DateTimeEncoder(json.JSONEncoder):
    def default(self, obj):
        if isinstance(obj, (datetime.date, datetime.datetime)):
            return obj.isoformat()
        return super(DateTimeEncoder, self).default(obj)

# 定义指数智能体状态
class IndexAgentState(BaseAgentState):
    symbol: str = ""
    analysis_type: str = ""
    index_data: Dict = {}
    # 各子分析结果
    basic_info_analysis: str = ""
    trend_analysis: str = ""
    component_analysis: str = ""
    technical_analysis: str = ""

class IndexAnalysisAgent(BaseAgent):
    """指数分析智能体"""
    
    def _build_workflow(self) -> StateGraph:
        """构建智能体工作流"""
        workflow = StateGraph(IndexAgentState)
        
        # 添加节点
        workflow.add_node("parse_query", self._parse_query)
        workflow.add_node("retrieve_data", self._retrieve_data)
        
        # 添加各个分析节点
        workflow.add_node("analyze_basic_info", self._analyze_basic_info)
        workflow.add_node("analyze_trend", self._analyze_trend)
        workflow.add_node("analyze_component", self._analyze_component)
        workflow.add_node("analyze_technical", self._analyze_technical)
        
        # 路由节点
        workflow.add_node("route_next", self._route_to_next_analysis)
        
        # 综合分析节点
        workflow.add_node("synthesize_analysis", self._synthesize_analysis)
        
        # 设置节点连接关系 - 基本工作流
        workflow.set_entry_point("parse_query")
        workflow.add_edge("parse_query", "retrieve_data")
        
        # 从数据获取到各个分析节点的连接
        workflow.add_edge("retrieve_data", "analyze_basic_info")
        
        # 从各个分析节点到路由节点的连接
        workflow.add_edge("analyze_basic_info", "route_next")
        workflow.add_edge("analyze_trend", "route_next")
        workflow.add_edge("analyze_component", "route_next")
        workflow.add_edge("analyze_technical", "route_next")
        
        # 使用条件边缘路由到下一个节点
        workflow.add_conditional_edges(
            "route_next",
            self._decide_next_node,
            {
                "analyze_trend": "analyze_trend",
                "analyze_component": "analyze_component",
                "analyze_technical": "analyze_technical",
                "synthesize_analysis": "synthesize_analysis"
            }
        )
        
        # 从综合分析到结束的连接
        workflow.add_edge("synthesize_analysis", END)
        
        return workflow
    
    def _decide_next_node(self, state: IndexAgentState) -> str:
        """确定下一个要执行的节点（为条件边缘决策函数）"""
        # 检查数据和分析类型
        index_data = state.get('index_data', {})
        analysis_type = state.get('analysis_type', '')
        
        # 趋势分析
        if analysis_type in ['趋势分析', '综合分析'] and 'history_data' in index_data and 'trend_analysis' not in state:
            return "analyze_trend"
        
        # 成分股分析
        if analysis_type in ['成分股分析', '综合分析'] and 'component_data' in index_data and 'component_analysis' not in state:
            return "analyze_component"
        
        # 技术分析
        if analysis_type in ['技术分析', '综合分析'] and 'history_data' in index_data and 'technical_analysis' not in state:
            return "analyze_technical"
        
        # 所有必要分析都已完成，进行综合分析
        return "synthesize_analysis"
    
    def _route_to_analysis_nodes(self, state: IndexAgentState) -> str:
        """根据分析类型和数据可用性路由到适当的分析节点"""
        # 首先始终进行基本信息分析
        return "analyze_basic_info"
    
    def _route_to_next_analysis(self, state: IndexAgentState) -> Dict:
        """检查并决定下一个分析节点 - 此函数仅用于边缘连接，不决定实际路由"""
        # 此函数作为节点函数，必须返回字典以更新状态
        # 在使用条件边缘时，节点的输出不决定路由，而是由条件边缘的决策函数决定
        return state
    
    def _parse_query(self, state: IndexAgentState) -> Dict:
        """解析用户查询并提取指数代码和分析类型"""
        try:
            # 为指数分析创建一个临时解析提示
            prompt = """
            从以下指数分析请求中提取信息：
            {user_query}
            
            只输出JSON格式，不要任何解释，格式为：
            {{"symbol": "指数代码", "analysis_type": "趋势分析|成分股分析|技术分析|综合分析"}}
            
            确保：
            1. 使用双引号
            2. 无注释
            """
            
            response = self.llm.invoke(prompt.format(user_query=state['user_query']))
            
            # 解析JSON
            result = self._clean_json_response(response.content)
            
            # 数据校验
            result['symbol'] = str(result.get('symbol', '000001')).strip()
            analysis_type = str(result.get('analysis_type', '趋势分析')).lower()
            result['analysis_type'] = {
                'trend': '趋势分析',
                'component': '成分股分析',
                'technical': '技术分析',
                'comprehensive': '综合分析',
                '趋势': '趋势分析',
                '成分股': '成分股分析',
                '技术': '技术分析',
                '综合': '综合分析'
            }.get(analysis_type, '趋势分析')
            
            return {
                "symbol": result['symbol'],
                "analysis_type": result['analysis_type'],
                "parsed_query": json.dumps(result, ensure_ascii=False, cls=DateTimeEncoder),
                "user_query": state['user_query']
            }
                
        except Exception as e:
            self.logger.error(f"查询解析失败: {str(e)}")
            # 默认返回
            return {
                "symbol": "000001",
                "analysis_type": "趋势分析",
                "parsed_query": f"解析失败：{str(e)}",
                "user_query": state['user_query']
            }
    
    def _retrieve_data(self, state: IndexAgentState) -> Dict:
        """获取相关数据"""
        try:
            symbol = state['symbol']
            analysis_type = state['analysis_type']
            
            # 获取指数对象
            index = get_asset(symbol, ASSET_TYPE_INDEX)
            
            # 基础数据 - 所有分析类型都需要
            basic_info = index.basic_info
            
            # 根据分析类型获取不同数据
            if analysis_type == '趋势分析':
                data = {
                    "basic_info": basic_info,
                    "history_data": index.get_history(limit=90).to_dict('records')
                }
            elif analysis_type == '成分股分析':
                data = {
                    "basic_info": basic_info,
                    "component_data": index.get_components()
                }
            elif analysis_type == '技术分析':
                data = {
                    "basic_info": basic_info,
                    "history_data": index.get_history(limit=120).to_dict('records')
                }
            elif analysis_type == '综合分析':
                data = {
                    "basic_info": basic_info,
                    "history_data": index.get_history(limit=90).to_dict('records'),
                    "component_data": index.get_components(),
                    "industry_data": index.get_industry_distribution()
                }
            else:
                data = {"error": f"暂不支持的分析类型: {analysis_type}"}
            
            return {
                **state,  # 保留现有状态
                "index_data": data
            }
                
        except Exception as e:
            self.logger.error(f"数据获取失败: {str(e)}")
            return {
                **state,
                "index_data": {"error": f"数据获取失败: {str(e)}"}
            }
    
    def _analyze_basic_info(self, state: IndexAgentState) -> Dict:
        """分析基本信息"""
        try:
            # 获取相关数据
            index_data = state.get('index_data', {})
            basic_info = index_data.get('basic_info', {})
            
            # 检查数据
            if not basic_info:
                return {**state, "basic_info_analysis": "未获取到基本信息数据"}
            
            # 为指数基本信息分析创建一个临时提示
            prompt = """
            请基于以下指数基本信息数据进行分析：

            基本信息：{basic_info}

            请生成一份简洁的指数基本信息分析总结，包含：
            1. 指数概述与特征
            2. 编制方法与规则
            3. 应用场景与代表性
            4. 关键特点分析

            使用Markdown格式，确保分析简洁明了但专业深入，不超过200字。
            """
            
            # 构建提示
            prompt_formatted = prompt.format(
                basic_info=json.dumps(basic_info, ensure_ascii=False, cls=DateTimeEncoder)
            )
            
            # 调用大模型
            response = self.llm.invoke(prompt_formatted)
            
            # 返回结果
            return {**state, "basic_info_analysis": response.content}
                
        except Exception as e:
            self.logger.error(f"基本信息分析失败: {str(e)}")
            return {**state, "basic_info_analysis": f"分析失败: {str(e)}"}
    
    def _analyze_trend(self, state: IndexAgentState) -> Dict:
        """分析趋势数据"""
        try:
            # 获取相关数据
            index_data = state.get('index_data', {})
            basic_info = index_data.get('basic_info', {})
            history_data = index_data.get('history_data', [])
            
            # 检查数据
            if not history_data:
                return {**state, "trend_analysis": "未获取到历史数据"}
            
            # 限制数据量 - 只取前30条
            limited_history = history_data[:30] if len(history_data) > 30 else history_data
            
            # 为指数趋势分析创建一个临时提示
            prompt = """
            请基于以下指数历史数据进行趋势分析：

            基本信息：{basic_info}
            历史数据：{history_data}

            请生成一份指数趋势分析总结，专注于：
            1. 长期趋势判断
            2. 中期走势分析
            3. 短期波动特征
            4. 关键转折点识别
            5. 趋势持续性评估

            使用Markdown格式，确保分析简洁明了但专业深入，不超过300字。
            """
            
            # 构建提示
            prompt_formatted = prompt.format(
                basic_info=json.dumps(basic_info, ensure_ascii=False, cls=DateTimeEncoder),
                history_data=json.dumps(limited_history, ensure_ascii=False, cls=DateTimeEncoder)
            )
            
            # 调用大模型
            response = self.llm.invoke(prompt_formatted)
            
            # 返回结果
            return {**state, "trend_analysis": response.content}
                
        except Exception as e:
            self.logger.error(f"趋势分析失败: {str(e)}")
            return {**state, "trend_analysis": f"分析失败: {str(e)}"}
    
    def _analyze_component(self, state: IndexAgentState) -> Dict:
        """分析成分股数据"""
        try:
            # 获取相关数据
            index_data = state.get('index_data', {})
            basic_info = index_data.get('basic_info', {})
            component_data = index_data.get('component_data', {})
            
            # 检查数据
            if not component_data:
                return {**state, "component_analysis": "未获取到成分股数据"}
            
            # 为指数成分股分析创建一个临时提示
            prompt = """
            请基于以下指数成分股数据进行分析：

            基本信息：{basic_info}
            成分股数据：{component_data}

            请生成一份成分股分析总结，专注于：
            1. 主要成分股权重分布
            2. 行业分布特征
            3. 重点成分股表现
            4. 成分股集中度分析
            5. 成分股波动性比较

            使用Markdown格式，确保分析简洁明了但专业深入，不超过300字。
            """
            
            # 构建提示
            prompt_formatted = prompt.format(
                basic_info=json.dumps(basic_info, ensure_ascii=False, cls=DateTimeEncoder),
                component_data=json.dumps(component_data, ensure_ascii=False, cls=DateTimeEncoder)
            )
            
            # 调用大模型
            response = self.llm.invoke(prompt_formatted)
            
            # 返回结果
            return {**state, "component_analysis": response.content}
                
        except Exception as e:
            self.logger.error(f"成分股分析失败: {str(e)}")
            return {**state, "component_analysis": f"分析失败: {str(e)}"}
    
    def _analyze_technical(self, state: IndexAgentState) -> Dict:
        """分析技术指标"""
        try:
            # 获取相关数据
            index_data = state.get('index_data', {})
            basic_info = index_data.get('basic_info', {})
            history_data = index_data.get('history_data', [])
            
            # 检查数据
            if not history_data:
                return {**state, "technical_analysis": "未获取到历史数据，无法计算技术指标"}
            
            # 限制数据量 - 只取前30条用于计算技术指标
            limited_history = history_data[:30] if len(history_data) > 30 else history_data
            
            # 为指数技术分析创建一个临时提示
            prompt = """
            请基于以下指数历史数据进行技术分析：

            基本信息：{basic_info}
            历史数据：{history_data}

            请计算并分析以下技术指标：
            1. 移动平均线（MA5、MA10、MA20、MA60）
            2. MACD指标
            3. KDJ指标
            4. RSI指标
            5. 布林带

            并生成一份技术指标分析总结，专注于各指标反映的市场状况和可能的信号。
            使用Markdown格式，确保分析简洁明了但专业深入，不超过300字。
            """
            
            # 构建提示
            prompt_formatted = prompt.format(
                basic_info=json.dumps(basic_info, ensure_ascii=False, cls=DateTimeEncoder),
                history_data=json.dumps(limited_history, ensure_ascii=False, cls=DateTimeEncoder)
            )
            
            # 调用大模型
            response = self.llm.invoke(prompt_formatted)
            
            # 返回结果
            return {**state, "technical_analysis": response.content}
                
        except Exception as e:
            self.logger.error(f"技术指标分析失败: {str(e)}")
            return {**state, "technical_analysis": f"分析失败: {str(e)}"}
    
    def _synthesize_analysis(self, state: IndexAgentState) -> Dict:
        """综合各个分析，生成最终报告"""
        try:
            # 收集所有分析总结
            summaries = {}
            for key in [
                "basic_info_analysis", "trend_analysis", "component_analysis", "technical_analysis"
            ]:
                if key in state and state.get(key):
                    summaries[key] = state.get(key)
            
            # 为指数综合分析创建一个临时提示
            prompt = """
            请基于以下各方面的分析总结，生成一份指数综合分析报告：

            用户需求：{user_query}
            指数代码：{symbol}

            各项分析总结：
            {summaries}

            请生成一份全面的指数综合分析报告，包含以下要素：
            1. 核心结论摘要
            2. 趋势分析要点
            3. 成分股结构分析要点
            4. 技术面分析要点
            5. 市场情绪评估
            6. 风险因素分析
            7. 投资策略建议

            使用Markdown格式输出，确保专业、客观、有深度，突出关键发现和投资建议。
            """
            
            # 构建提示
            prompt_formatted = prompt.format(
                user_query=state.get('user_query', ''),
                symbol=state.get('symbol', ''),
                summaries=json.dumps(summaries, ensure_ascii=False)
            )
            
            # 调用大模型
            response = self.llm.invoke(prompt_formatted)
            
            # 返回结果
            return {**state, "analysis_result": response.content}
                
        except Exception as e:
            self.logger.error(f"综合分析失败: {str(e)}")
            return {**state, "analysis_result": self._format_error_response(str(e))}
    
    def analyze(self, query: str) -> str:
        """执行完整分析流程"""
        try:
            result = self.agent.invoke({"user_query": query})
            return result['analysis_result']
        except Exception as e:
            self.logger.error(f"分析执行失败: {str(e)}")
            return self._format_error_response(str(e))